#pragma once
namespace DEVICE 
{

template<typename T>
inline void gemm(cublasHandle_t handle, T* A,  T* B, T* C, int m, int n, int k, int lda, int ldb, int ldc, T alpha,  T beta) {
    if constexpr (std::is_same_v<T, float>) {
        cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
    } else if constexpr (std::is_same_v<T, double>) {
        cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}


template<typename T>
inline void trmm(cublasHandle_t handle, char _side, char _uplo, char _diag, const T* A, const T* B, T* C, int m, int n, int lda, int ldb, int ldc, T alpha) {
    assert(_side == 'L' || _side == 'R');
    assert(_uplo == 'L' || _uplo == 'U');
    assert(_diag == 'U' || _diag == 'N');
    cublasSideMode_t side = _side == 'L' ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
    cublasFillMode_t uplo = _uplo == 'L' ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
    cublasDiagType_t diag = _diag == 'U' ? CUBLAS_DIAG_UNIT : CUBLAS_DIAG_NON_UNIT;
    if constexpr (std::is_same_v<T, float>) {
        cublasStrmm(handle, side, uplo, CUBLAS_OP_N, diag, m, n, &alpha, A, lda, B, ldb, C, ldc);
    } else if constexpr (std::is_same_v<T, double>) {
        cublasDtrmm(handle, side, uplo, CUBLAS_OP_N, diag, m, n, &alpha, A, lda, B, ldb, C, ldc);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}

template<typename T>
inline void trsm(cublasHandle_t handle, char _side, char _uplo, char _diag, const T* A, T* B, int m, int n, int lda, int ldb) {
    assert(_side == 'L' || _side == 'R');
    assert(_uplo == 'L' || _uplo == 'U');
    assert(_diag == 'U' || _diag == 'N');
    cublasSideMode_t side = _side == 'L' ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
    cublasFillMode_t uplo = _uplo == 'L' ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
    cublasDiagType_t diag = _diag == 'U' ? CUBLAS_DIAG_UNIT : CUBLAS_DIAG_NON_UNIT;
    const T alpha = 1.0;
    if constexpr (std::is_same_v<T, float>) {
        cublasStrsm(handle, side, uplo, CUBLAS_OP_N, diag, m, n, &alpha, A, lda, B, ldb);
    } else if constexpr (std::is_same_v<T, double>) {
        cublasDtrsm(handle, side, uplo, CUBLAS_OP_N, diag, m, n, &alpha, A, lda, B, ldb);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}

template<typename T>
inline void getrf(cusolverDnHandle_t handle, T* workspace, T* A,  int m, int n, int lda,  int *devInfo) { 
    if constexpr (std::is_same_v<T, float>) {
        cusolverDnSgetrf(handle, m, n, A, lda, workspace, NULL, devInfo );
    } else if constexpr (std::is_same_v<T, double>) {
        cusolverDnDgetrf(handle, m, n, A, lda, workspace, NULL, devInfo );
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}

template<typename T>
inline void getrs(cusolverDnHandle_t handle, const T* A, const int *devIpiv, T* B, int n, int nrhs, int lda, int ldb, int *devInfo) {
    if constexpr (std::is_same_v<T, float>) {
        cusolverDnSgetrs(handle, CUBLAS_OP_N,  n, nrhs, A, lda, devIpiv, B, ldb, devInfo);
    } else if constexpr (std::is_same_v<T, double>) {
        cusolverDnDgetrs(handle, CUBLAS_OP_N,  n, nrhs, A, lda, devIpiv, B, ldb, devInfo);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}

// template<typename T>
// inline void getri(cublasHandle_t handle, const T* matrix, T* inverse, int n, int lda, int ldb, int *devInfo) {
//     const T* const Aarray[] = {matrix};
//     T* Carray[] = {inverse};
//     int *infoArray = {devInfo};
//     if constexpr (std::is_same_v<T, float>) {
//         cublasSgetriBatched(handle, n, Aarray, lda, NULL, Carray, ldb, infoArray, 1);
//     } else if constexpr (std::is_same_v<T, double>) {
//         cublasDgetriBatched(handle, n, Aarray, lda, NULL, Carray, ldb, infoArray, 1);
//     } else {
//         // Handle other types or provide an error message
//         static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
//     }
// }

template<typename T>
inline void gemv(cublasHandle_t handle, T* result, const T* matrix, const T* vec, int m, int n, int lda, T alpha, T beta) {
    if constexpr (std::is_same_v<T, float>) {
        cublasSgemv(handle, CUBLAS_OP_N, m, n, &alpha, matrix, lda, vec, 1, &beta, result, 1);
    } else if constexpr (std::is_same_v<T, double>) {
        if(fabs(alpha - 1.0) < 1e-12 && fabs(beta) < 1e-12 && lda % 2 == 0)
            new_dgemv(matrix, vec, result, n, lda, n, n);
        else
            cublasDgemv(handle, CUBLAS_OP_N, m, n, &alpha, matrix, lda, vec, 1, &beta, result, 1);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}

template<typename T>
inline void trmv(cublasHandle_t handle, char _uplo, char _diag, const T* matrix, T* vec, int n, int lda) {
    assert(_uplo == 'L' || _uplo == 'U');
    assert(_diag == 'U' || _diag == 'N');
    cublasFillMode_t uplo = _uplo == 'L' ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
    cublasDiagType_t diag = _diag == 'U' ? CUBLAS_DIAG_UNIT : CUBLAS_DIAG_NON_UNIT;
    if constexpr (std::is_same_v<T, float>) {
        cublasStrmv(handle, uplo, CUBLAS_OP_N, diag, n, matrix, lda, vec, 1);
    } else if constexpr (std::is_same_v<T, double>) {
        cublasDtrmv(handle, uplo, CUBLAS_OP_N, diag, n, matrix, lda, vec, 1);
    } else {
        // Handle other types or provide an error message
        static_assert(std::is_same_v<T, float> || std::is_same_v<T, double> , "Unsupported type");
    }
}
};
